from django.core.management.base import BaseCommand
from stocksys import models as stocksys_models


class Command(BaseCommand):
    help = 'create_stock_code_data'

    def get_stock_code_dir(self, path, first_name):
        with open(path) as file:
            lines = file.readlines()
            lines = [line.replace("\n", "").split("(") for line in lines]
            return [{"name": line[0], "search_code": first_name + line[1][:-1], "code": line[1][:-1]} for line in lines]

    def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('初始化数据'))
        self.stdout.write(self.style.SUCCESS('载入股票代码数据'))

        index_sh_id = stocksys_models.StockClass.objects.get(stock_class_name="sh").stock_class_id
        index_sz_id = stocksys_models.StockClass.objects.get(stock_class_name="sz").stock_class_id

        code_sz = self.get_stock_code_dir("./data/code_sz.txt", "sz")
        code_sh = self.get_stock_code_dir("./data/code_sh.txt", "sh")

        self.stdout.write(self.style.SUCCESS('导入深圳股票代码数据'))
        [stocksys_models.Stock.objects.create(stock_class=index_sz_id,
                                              stock_name=code_sz_line['name'],
                                              stock_api_name="sz" + code_sz_line['code'],
                                              stock_code=code_sz_line['code']
                                              ) for code_sz_line in code_sz]

        self.stdout.write(self.style.SUCCESS('导入上海股票代码数据'))
        [stocksys_models.Stock.objects.create(stock_class=index_sh_id,
                                              stock_name=code_sh_line['name'],
                                              stock_api_name="sh" + code_sh_line['code'],
                                              stock_code=code_sh_line['code']) for code_sh_line in code_sh]
        self.stdout.write(self.style.SUCCESS('结束载入'))
